# https://zhuanlan.zhihu.com/p/630200571
# 1 LLM model
from typing import Any, List
from transformers import AutoTokenizer, AutoModel


class chatGLM():
    def __init__(self, model_name, quantization_bit=4) -> None:
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        model = AutoModel.from_pretrained(model_name, trust_remote_code=True).half().cuda().eval()
        self.model = model.quantize(quantization_bit)

    def __call__(self, prompt) -> Any:
        response, _ = self.model.chat(self.tokenizer, prompt)  # 这里演示未使用流式接口. stream_chat()
        return response


llm = chatGLM(model_name="THUDM/chatglm-6b")
prompt = "你好"
response = llm(prompt)
print("response: %s" % response)

# ----------------------------------------------------------
# 2 使用prompt模板，格式化生成新的prompt
from langchain import PromptTemplate

template = "请给我解释一下{query}的意思"
promptTem = PromptTemplate(input_variables=["query"], template=template)
prompt = promptTem.format(query="天道酬勤")

# ----------------------------------------------------------
# 3 用chain连链接llm和prompt组件
from langchain.chains.base import Chain


class DemoChain():
    def __init__(self, llm, prompt) -> None:
        self.llm = llm
        self.prompt = prompt

    def run(self, query, context=None) -> Any:
        if context is not None:
            prompt = self.prompt.format(query=query, context=context)
        else:
            prompt = self.prompt.format(query=query)
        print("query=%s  -> prompt=%s" % (query, prompt))
        print("*" * 60)
        response = self.llm(prompt)
        return response


chain = DemoChain(llm=llm, prompt=promptTem)
print("-" * 80)
chain.run(query="天道酬勤")
print("-" * 80)

# ----------------------------------------------------------
# 4 示例 Embedding 和 vs-store
from langchain.embeddings.huggingface import HuggingFaceEmbeddings

embeddings = HuggingFaceEmbeddings(model_name="GanymedeNil_text2vec-large-chinese",
                                   model_kwargs={'device': "cuda"})
query_result = embeddings.embed_query("天道酬勤")
print("embedding query.shape=", np.array(query_result).shape)

texts = """ '天道酬勤'并不是鼓励人们不劳而获，而是提醒人们要遵循自然规律，通过不断的努力和付出来追求自己的目标。\n这种努力不仅仅是指身体上的劳动，
也包括精神上的努力和思考，以及学习和适应变化的能力。\n只要一个人具备这些能力，他就有可能会获得成功。"""

from langchain.text_splitter import CharacterTextSplitter
from langchain.docstore.document import Document


class TextSpliter(CharacterTextSplitter):
    def __init__(self, separator: str = "\n\n", **kwargs: Any):
        super().__init__(separator, **kwargs)

    def split_text(self, text: str) -> List[str]:
        texts = text.split("\n")
        texts = [Document(page_content=text, metadata={"from": "filename or book.txt"}) for text in texts]
        return texts


text_splitter = TextSpliter()
texts = text_splitter.split_text(texts)
texts1 = [text.page_content for text in texts]

vs_path = "demo-vs"
from langchain.vectorstores import FAISS

docs = embeddings.embed_documents(texts1)
vector_store = FAISS.from_documents(texts, embeddings)
vector_store.save_local(vs_path)

vector_store = FAISS.load_local(vs_path, embeddings)
related_docs_with_score = vector_store.similarity_search_with_score(query="天道酬勤", k=2)

# -------------
# 5 基于查询到的知识做prompt
context = ""
for pack in related_docs_with_score:
    doc, socre = pack
    content = doc.page_content
    print("检索到的知识=%s, from=%s, socre=%.3f" % (content, doc.metadata.get("from"), socre))
    context += content

# 重新配置一个基于上下文的模板在来调下语言模型
template = "已知{context}, 请给我解释一下{query}的意思?"
promptTem = PromptTemplate(input_variables=["context", "query"], template=template)
chain = DemoChain(llm=llm, prompt=promptTem)
print("-" * 80)
print(chain.run(query="天道酬勤", context=context))
print("-" * 80)

'''
----------------------log
response: 你好！请问有什么需要帮助的吗？
--------------------------------------------------------------------------------
query = 天道酬勤  -> prompt = 请给我解释一下天道酬勤的意思
--------------------------------------------------------------------------------
embedding
query.shape = (1024,)
检索到的知识 = 天道酬勤”并不是鼓励人们不劳而获，而是提醒人们要遵循自然规律，通过不断的努力和付出来追求自己的目标。, from=filename or book.txt, socre = 373.131
检索到的知识 = 这种努力不仅仅是指身体上的劳动，, from=filename or book.txt, socre = 740.042
--------------------------------------------------------------------------------
query = 天道酬勤  -> prompt = 已知天道酬勤”并不是鼓励人们不劳而获，而是提醒人们要遵循自然规律，通过不断的努力和付出来追求自己的目标。这种努力不仅仅是指身体上的劳动，, 请给我解释一下天道酬勤的意思?
“天道酬勤”是指遵循自然规律，通过不断的努力和付出来追求自己的目标，这种努力不仅仅是指身体上的劳动，还包括精神上的、心理上的和智力上的。它鼓励人们不要放弃，坚持追求自己的目标，即使需要付出很大的努力和时间。最终，这种努力和付出会转化为回报，使人们的生活更加充实和有意义。

这句话的意义在于提醒人们要坚持不懈地追求自己的目标，不要懒惰和放弃。同时，它也提醒人们要认识到自然规律，并遵循这些规律，才能取得成功。这三个方面的努力和付出是“天道酬勤”的核心意思。
'''